{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Video processing, Object detection & Tracking\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "**Demonstrating the video processing capabilities of Stone Soup**\n\nThis notebook will guide you progressively through the steps necessary to:\n\n1. Use the Stone Soup :class:`~.FrameReader` components to open and process video data;\n2. Use the :class:`~.TensorFlowBoxObjectDetector` to detect objects in video data, making use of TensorFlow object detection models;\n3. Build a :class:`~.MultiTargetTracker` to perform tracking of multiple object in video data.\n\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Software dependencies\nBefore we begin with this tutorial, there are a few things that we need to install in order to\nproceed.\n\n### FFmpeg\nFFmpeg is a free and open-source project consisting of a vast software suite of libraries and\nprograms for handling video, audio, and other multimedia files and streams. Stone Soup (or more\naccurately some of its extra dependencies) make use of FFmpeg to read and output video. Download\nlinks and installation instructions for FFmpeg can be found `here <https://www.ffmpeg.org/download.html>`__.\n\n### TensorFlow\nTensorFlow is a free and open-source software library for dataflow and differentiable programming\nacross a range of tasks, such machine learning. TensorFlow includes an Object Detection API that\nmakes it easy to construct, train and deploy object detection models, as well as a collection of\npre-trained models that can be used for out-of-the-box inference. A quick TensorFlow installation\ntutorial can be found `here <https://tensorflow-object-detection-api-tutorial.readthedocs.io/en/latest/install.html>`__.\n\n### Stone Soup\nTo perform video-processing using Stone Soup, we need to install some extra dependencies. The\neasiest way to achieve this is by running the following commands in a Terminal window:\n\n.. code::\n\n    git clone \"https://github.com/dstl/Stone-Soup.git\"\n    cd Stone-Soup\n    python -m pip install -e .[dev,video,tensorflow]\n\n### Pytube\nWe will also use pytube_ to download a Youtube video for the purposes of this tutorial. In the\nsame Terminal window, run the following command to install ``pytube``:\n\n.. code::\n\n    pip install pytube\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Using the Stone Soup :class:`~.FrameReader` classes\nThe :class:`~.FrameReader` abstract class is intended as the base class for Stone Soup readers\nthat read frames from any form of imagery data. As of now, Stone Soup has two implementations of\n:class:`~.FrameReader` subclasses:\n\n1. The :class:`~.VideoClipReader` component, which uses MoviePy_ to read video frames from a file.\n2. The :class:`~.FFmpegVideoStreamReader` component, which uses ffmpeg-python_ to read frames from real-time video streams (e.g. RTSP).\n\nIn this tutorial we will focus on the :class:`~.VideoClipReader`, since setting up a stream for\nthe :class:`~.FFmpegVideoStreamReader` is more involved. Nevertheless, the use and interface of\nthe two readers is mostly identical after initialisation and an example of how to initialise the\nlater will also be provided\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Download and store the video\nFirst we will download the video that we will use throughout this tutorial. The code snippet\nshown below will download the video and save it your working directory as ``sample1.mp4``.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import os\nfrom pytube import YouTube\nVIDEO_FILENAME = 'sample1'\nVIDEO_EXTENTION = '.mp4'\nVIDEO_PATH = os.path.join(os.getcwd(), VIDEO_FILENAME+VIDEO_EXTENTION)\n\nif not os.path.exists(VIDEO_PATH):\n    yt = YouTube('http://www.youtube.com/watch?v=MNn9qKG2UFI')\n    yt.streams.get_by_itag(18).download(filename=VIDEO_PATH)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Building the video reader\n\n#### VideoClipReader\nWe will use the :class:`~.VideoClipReader` class to read and replay the downloaded file. We also\nconfigure the reader to only replay the clip for the a duration of 2 seconds between `00:10` and\n`00:12`.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import datetime\nfrom stonesoup.reader.video import VideoClipReader\nstart_time = datetime.timedelta(minutes=0, seconds=10)\nend_time = datetime.timedelta(minutes=0, seconds=12)\nframe_reader = VideoClipReader(VIDEO_PATH, start_time, end_time)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "It is also possible to apply clip transformations and effects, as per the\n`MoviePy documentation <https://zulko.github.io/moviepy/getting_started/effects.html>`_.\nThe underlying MoviePy :class:`~VideoFileClip` instance can be accessed through the\n:attr:`~.VideoClipReader.clip` class property. For example, we can crop out 100 pixels from\nthe top and left of the frames, as they are read by the reader, as shown below.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from moviepy.video.fx import all\nframe_reader.clip = all.crop(frame_reader.clip, 100, 100)\nnum_frames = len(list(frame_reader.clip.iter_frames()))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### FFmpegVideoStreamReader\nFor reference purposes, we also include here an example of how to build a\n:class:`~.FFmpegVideoStreamReader`. Let's assume that we have a camera which broadcasts its feed\nthrough a public RTSP stream, under the URL ``rtsp://192.168.55.10:554/stream``. We can build a\n:class:`~.FFmpegVideoStreamReader` object to read frames from this stream as follows:\n\n.. code:: python\n\n  in_opts = {'threads': 1, 'fflags': 'nobuffer'}\n  out_opts = {'format': 'rawvideo', 'pix_fmt': 'bgr24'}\n  stream_url = 'rtsp://192.168.55.10:554/stream'\n  video_reader = FFmpegVideoStreamReader(stream_url, input_opts=in_opts, output_opts=out_opts)\n\n.. important::\n\n  Note that the above code is an illustrative example and will not be run.\n\n:attr:`~.FFmpegVideoStreamReader.input_opts` and :attr:`~.FFmpegVideoStreamReader.output_opts`\nare optional arguments, which allow users to specify options for the input and output FFmpeg\nstreams, as documented by `FFmpeg <https://ffmpeg.org/ffmpeg.html#toc-Options>`__ and\nffmpeg-python_.\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Reading frames from the reader\nAll :class:`~.FrameReader` objects, of which the :class:`~.VideoClipReader` is a subclass,\ngenerate frames in the form of :class:`~.ImageFrame` objects. Below we show an example of how to\nread and visualise these frames using `matplotlib`.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from copy import copy\nfrom PIL import Image\nfrom matplotlib import pyplot as plt\nfrom matplotlib import animation\n\nfig, ax = plt.subplots(num=\"VideoClipReader output\")\nartists = []\n\nprint('Running FrameReader example...')\nfor timestamp, frame in frame_reader:\n    if not (len(artists)+1) % 10:\n        print(\"Frame: {}/{}\".format(len(artists)+1, num_frames))\n\n    # Read the frame pixels\n    pixels = copy(frame.pixels)\n\n    # Plot output\n    image = Image.fromarray(pixels)\n    ax.axes.xaxis.set_visible(False)\n    ax.axes.yaxis.set_visible(False)\n    fig.tight_layout()\n    artist = ax.imshow(image, animated=True)\n    artists.append([artist])\n\nani = animation.ArtistAnimation(fig, artists, interval=20, blit=True, repeat_delay=200)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Using the :class:`~.TensorFlowBoxObjectDetector` class\nWe now continue by demonstrating how to use the :class:`~.TensorFlowBoxObjectDetector` to detect\nobjects, and more specifically cars, within the frames read in by our ``video_reader``. The\n:class:`~.TensorFlowBoxObjectDetector` can utilise both pre-trained and custom-trained TensorFlow\nobject detection models which generate detection in the form of bounding boxes. In this example,\nwe will make use of a pre-trained model from the\n`TensorFlow detection model zoo <https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/tf2_detection_zoo.md>`_,\nbut the process of using a custom-trained TensorFlow model is the same.\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Downloading the model\nThe code snippet shown below is used to download the object detection model that we will feed\ninto the :class:`~.TensorFlowBoxObjectDetector`, as well as the label file (.pbtxt) which\ncontains a list of strings used to add the correct label to each detection (e.g. car).\n\nThe particular detection algorithm we will use is the Faster-RCNN, with an Inception\nResnet v2 backbone and running in Atrous mode with low proposals, pre-trained on the MSCOCO\ndataset.\n\n<div class=\"alert alert-danger\"><h4>Warning</h4><p>**The downloaded model has a size of approximately 500 MB**. Therefore it is advised that you\n  run the script on a stable (ideally not mobile) internet connection. The files will only be\n  downloaded the first time the script is run. In consecutive runs the code will skip this step,\n  provided that ``PATH_TO_MODEL`` and ``PATH_TO_LABELS`` are valid paths.</p></div>\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import os\nos.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'    # Suppress TensorFlow logging (1)\nimport pathlib\nimport tensorflow as tf\n\ntf.get_logger().setLevel('ERROR')           # Suppress TensorFlow logging (2)\n\n# Enable GPU dynamic memory allocation\ngpus = tf.config.experimental.list_physical_devices('GPU')\nfor gpu in gpus:\n    tf.config.experimental.set_memory_growth(gpu, True)\n\n# Download and extract model\ndef download_model(model_name):\n    base_url = 'http://download.tensorflow.org/models/object_detection/'\n    model_file = model_name + '.tar.gz'\n    model_dir = tf.keras.utils.get_file(fname=model_name,\n                                        origin=base_url + model_file,\n                                        untar=True)\n    model_dir = pathlib.Path(model_dir)/\"saved_model\"\n    return str(model_dir)\n\nMODEL_NAME = 'faster_rcnn_inception_resnet_v2_atrous_lowproposals_coco_2018_01_28'\nPATH_TO_MODEL = download_model(MODEL_NAME)\n\n# Download labels file\ndef download_labels(filename):\n    base_url = 'https://raw.githubusercontent.com/tensorflow/models/master/research/object_detection/data/'\n    label_dir = tf.keras.utils.get_file(fname=filename,\n                                        origin=base_url + filename,\n                                        untar=False)\n    label_dir = pathlib.Path(label_dir)\n    return str(label_dir)\n\nLABEL_FILENAME = 'mscoco_label_map.pbtxt'\nPATH_TO_LABELS = download_labels(LABEL_FILENAME)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Building the detector\nNext, we proceed to initialise our detector object. To do this, we require the ``frame_reader``\nobject we built previously, as well as a path to the (downloaded) ``saved_model`` directory and\nlabel (.pbtxt) file, which we have already defined above under the ``PATH_TO_MODEL`` and\n``PATH_TO_LABELS`` variables.\n\nThe :class:`~.TensorFlowBoxObjectDetector` object can optionally be configured to digest frames\nfrom the provided reader asynchronously, and only perform detection on the last frame digested,\nby setting ``run_async=True``.This is suitable when the detector is applied to readers generating\na live feed (e.g. the :class:`~.FFmpegVideoStreamReader`), where real-time processing is\nparamount. Since we are using a :class:`~.VideoClipReader` in this example, we set\n``run_async=False``, which is also the default setting.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from stonesoup.detector.tensorflow import TensorFlowBoxObjectDetector\n\nrun_async = False                           # Configure the detector to run in synchronous mode\ndetector = TensorFlowBoxObjectDetector(frame_reader, PATH_TO_MODEL, PATH_TO_LABELS,\n                                       run_async=run_async)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Filtering-out unwanted detections\nIn this section we showcase how we can utilise Stone Soup :class:`~.Feeder` objects in order to\nfilter out unwanted detections. One example of feeder we can use is the\n:class:`~.MetadataValueFilter`, which allows us to filter detections by applying a custom\noperator on particular fields of the :attr:`~.Detection.metadata` property of detections.\n\nEach detection generated by :class:`~.TensorFlowBoxObjectDetector` carries the following\n:attr:`~.Detection.metadata` fields:\n\n - ``raw_box``: The raw bounding box containing the normalised coordinates ``[y_0, x_0, y_1, x_1]``, as generated by TensorFlow.\n - ``class``: A dict with keys ``id`` and ``name`` relating to the id and name of the detection class.\n - ``score``: A float in the range ``(0, 1]`` indicating the detector's confidence.\n\nDetection models trained on the MSCOCO dataset, such as the one we downloaded, are able to detect\n90 different classes of objects (see the `downloaded .pbtxt file <https://github.com/tensorflow/models/blob/master/research/object_detection/data/mscoco_label_map.pbtxt>`_\nfor a full list). Instead, as we discussed at the beginning of the tutorial, we wish to limit the\ndetections to only those classified as cars. This can be done as follows:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from stonesoup.feeder.filter import MetadataValueFilter\ndetector = MetadataValueFilter(detector, 'class', lambda x: x['name'] == 'car')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Continuing, we may want to filter out detections which have a low confidence score:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "detector = MetadataValueFilter(detector, 'score', lambda x: x > 0.1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Finally, we observed that the detector tends to incorrectly generate detections which are much\nlarger the the size we expect for a car. Therefore, we can filter out those detections by only\nallowing ones whose width is less the 20\\% of the frame width (i.e. ``x_1-x_0 < 0.2``):\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "detector = MetadataValueFilter(detector, 'raw_box', lambda x: x[3]-x[1] < 0.2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "You are encouraged to comment out any/all of the above filter definitions and observe the\nproduced output.\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Reading and visualising detections\nDetections generated by the :class:`~.TensorFlowBoxObjectDetector` have a 4-dimensional\n:attr:`~.Detection.state_vector` in the form of a bounding boxes that captures the area of the\nframe where an object is detected. Each bounding box is represented by a vector of the form\n``[x, y, w, h]``, where ``x, y`` denote the relative pixel coordinates of the top-left corner,\nwhile ``w, h`` denote the relative width and height of the bounding box. Below we show an example\nof how to read and visualise these detections using `matplotlib`.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import numpy as np\nfrom PIL import ImageDraw\n\n\ndef draw_detections(image, detections, show_class=False, show_score=False):\n    \"\"\" Draw detections on an image\n\n    Parameters\n    ----------\n    image: :class:`PIL.Image`\n        Image on which to draw the detections\n    detections: : set of :class:`~.Detection`\n        A set of detections generated by :class:`~.TensorFlowBoxObjectDetector`\n    show_class: bool\n        Whether to draw the class of the object. Default is ``False``\n    show_score: bool\n        Whether to draw the score of the object. Default is ``False``\n\n    Returns\n    -------\n    : :class:`PIL.Image`\n        Image with detections drawn\n    \"\"\"\n    draw = ImageDraw.Draw(image)\n    for detection in detections:\n        x0, y0, w, h = np.array(detection.state_vector).reshape(4)\n        x1, y1 = (x0 + w, y0 + h)\n        draw.rectangle([x0, y0, x1, y1], outline=(0, 255, 0), width=1)\n        class_ = detection.metadata['class']['name']\n        score = round(float(detection.metadata['score']),2)\n        if show_class and show_score:\n            draw.text((x0,y1 + 2), '{}:{}'.format(class_, score), fill=(0, 255, 0))\n        elif show_class:\n            draw.text((x0, y1 + 2), '{}'.format(class_), fill=(0, 255, 0))\n        elif show_score:\n            draw.text((x0, y1 + 2), '{}'.format(score), fill=(0, 255, 0))\n\n    del draw\n    return image\n\n\nfig2, ax2 = plt.subplots(num=\"TensorFlowBoxObjectDetector output\")\nartists2 = []\nprint(\"Running TensorFlowBoxObjectDetector example... Be patient...\")\nfor timestamp, detections in detector:\n    if not (len(artists2)+1) % 10:\n        print(\"Frame: {}/{}\".format(len(artists2)+1, num_frames))\n\n    # Read the frame pixels\n    frame = frame_reader.frame\n    pixels = copy(frame.pixels)\n\n    # Plot output\n    image = Image.fromarray(pixels)\n    image = draw_detections(image, detections, True, True)\n    ax2.axes.xaxis.set_visible(False)\n    ax2.axes.yaxis.set_visible(False)\n    fig2.tight_layout()\n    artist = ax2.imshow(image, animated=True)\n    artists2.append([artist])\n\nani2 = animation.ArtistAnimation(fig2, artists2, interval=20, blit=True, repeat_delay=200)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Constructing a Multi-Object Video Tracker\nIn this final segment of the tutorial we will see how we can use the above demonstrated\ncomponents to perform tracking of multiple objects within Stone Soup.\n\n### Defining the state-space models\nTransition Model\n****************\nWe begin our definition of the state-space models by defining the hidden state\n$\\mathrm{x}_k$, i.e. the state that we wish to estimate:\n\n\\begin{align}\\mathrm{x}_k = [x_k, \\dot{x}_k, y_k, \\dot{y}_k, w_k, h_k]\\end{align}\n\nwhere $x_k, y_k$ denote the pixel coordinates of the top-left corner of the bounding box\ncontaining an object, with $\\dot{x}_k, \\dot{y}_k$ denoting their respective rate of change,\nwhile $w_k$ and $h_k$ denote the width and height of the box, respectively.\n\nWe assume that $x_k$ and $y_k$ move with nearly :class:`~.ConstantVelocity`, while\n$w_k$ and $h_k$ evolve according to a :class:`~.RandomWalk`.Using these assumptions,\nwe proceed to construct our Stone Soup :class:`~.TransitionModel` as follows:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from stonesoup.models.transition.linear import (CombinedLinearGaussianTransitionModel,\n                                                ConstantVelocity, RandomWalk)\nt_models = [ConstantVelocity(20**2), ConstantVelocity(20**2), RandomWalk(20**2), RandomWalk(20**2)]\ntransition_model = CombinedLinearGaussianTransitionModel(t_models)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Measurement Model\nContinuing, we define the measurement state $\\mathrm{y}_k$, which follows naturally from\nthe form of the detections generated by the :class:`~.TensorFlowBoxObjectDetector` we previously\ndiscussed:\n\n\\begin{align}\\mathrm{y}_k = [x_k, y_k, w_k, h_k]\\end{align}\n\nWe make use of a 4-dimensional :class:`~.LinearGaussian` model as our :class:`~.MeasurementModel`,\nwhereby we can see that the individual indices of $\\mathrm{y}_k$ map to indices `[0,2,4,5]`\nof the 6-dimensional state $\\mathrm{x}_k$:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from stonesoup.models.measurement.linear import LinearGaussian\nmeasurement_model = LinearGaussian(ndim_state=6, mapping=[0, 2, 4, 5],\n                                   noise_covar=np.diag([1**2, 1**2, 3**2, 3**2]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Defining the tracker components\nWith the state-space models defined, we proceed to build our tracking components\n\n#### Filtering\nSince we have assumed Linear-Gaussian models, we will be using a Kalman Filter to perform\nfiltering of the underlying single-target densities. This is done by making use of the\n:class:`~.KalmanPredictor` and :class:`~.KalmanUpdater` classes, which we define below:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from stonesoup.predictor.kalman import KalmanPredictor\npredictor = KalmanPredictor(transition_model)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from stonesoup.updater.kalman import KalmanUpdater\nupdater = KalmanUpdater(measurement_model)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<div class=\"alert alert-info\"><h4>Note</h4><p>For more information on the above classes and how they operate you can refer to the Stone\n  Soup tutorial on\n  `using the Kalman Filter <https://stonesoup.readthedocs.io/en/latest/auto_tutorials/01_KalmanFilterTutorial.html>`_.</p></div>\n\n#### Data Association\nWe utilise a :class:`~.DistanceHypothesiser` to generate hypotheses between tracks and\nmeasurements, where :class:`~.Mahalanobis` distance is used as a measure of quality:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from stonesoup.hypothesiser.distance import DistanceHypothesiser\nfrom stonesoup.measures import Mahalanobis\nhypothesiser = DistanceHypothesiser(predictor, updater, Mahalanobis(), 10)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Continuing the :class:`~.GNNWith2DAssigment` class is used to perform fast joint data association,\nbased on the Global Nearest Neighbour (GNN) algorithm:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from stonesoup.dataassociator.neighbour import GNNWith2DAssignment\ndata_associator = GNNWith2DAssignment(hypothesiser)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<div class=\"alert alert-info\"><h4>Note</h4><p>For more information on the above classes and how they operate you can refer to the\n  `Data Association - clutter <https://stonesoup.readthedocs.io/en/latest/auto_tutorials/05_DataAssociation-Clutter.html>`_\n  and `Data Association - Multi-Target Tracking <https://stonesoup.readthedocs.io/en/latest/auto_tutorials/06_DataAssociation-MultiTargetTutorial.html>`_\n  tutorials.</p></div>\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Track Initiation\nFor initialising tracks we will use a :class:`~.MultiMeasurementInitiator`, which allows our\ntracker to tentatively initiate tracks from unassociated measurements, and hold them within the\ninitiator until they have survived for at least 10 frames. We also define a\n:class:`~.UpdateTimeStepsDeleter` deleter to be used by the initiator to delete tentative tracks\nthat have not been associated to a measurement in the last 3 frames.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from stonesoup.types.state import GaussianState\nfrom stonesoup.types.array import CovarianceMatrix, StateVector\nfrom stonesoup.initiator.simple import MultiMeasurementInitiator\nfrom stonesoup.deleter.time import UpdateTimeStepsDeleter\nprior_state = GaussianState(StateVector(np.zeros((6,1))),\n                            CovarianceMatrix(np.diag([100**2, 30**2, 100**2, 30**2, 100**2, 100**2])))\ndeleter_init = UpdateTimeStepsDeleter(time_steps_since_update=3)\ninitiator = MultiMeasurementInitiator(prior_state, deleter_init, data_associator, updater,\n                                      measurement_model, min_points=10)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Track Deletion\nFor confirmed tracks we used again a :class:`~.UpdateTimeStepsDeleter`, but this time configured\nto delete tracks after they have not bee associated to a measurement in the last 15 frames.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "deleter = UpdateTimeStepsDeleter(time_steps_since_update=15)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<div class=\"alert alert-info\"><h4>Note</h4><p>For more information on the above classes and how they operate you can refer to the Stone\n  `Initiators & Deleters <https://stonesoup.readthedocs.io/en/latest/auto_tutorials/09_Initiators_&_Deleters.html>`_\n  tutorial.</p></div>\n\n### Building the tracker\nNow that we have defined all our tracker components we proceed to build our multi-target tracker:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from stonesoup.tracker.simple import MultiTargetTracker\ntracker = MultiTargetTracker(\n    initiator=initiator,\n    deleter=deleter,\n    detector=detector,\n    data_associator=data_associator,\n    updater=updater,\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Running the tracker\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def draw_tracks(image, tracks, show_history=True, show_class=True, show_score=True):\n    \"\"\" Draw tracks on an image\n\n    Parameters\n    ----------\n    image: :class:`PIL.Image`\n        Image on which to draw the tracks\n    detections: : set of :class:`~.Tracks`\n        A set of tracks generated by our :class:`~.MultiTargetTracker`\n    show_history: bool\n        Whether to draw the trajectory of the track. Default is ``True``\n    show_class: bool\n        Whether to draw the class of the object. Default is ``True``\n    show_score: bool\n        Whether to draw the score of the object. Default is ``True``\n\n    Returns\n    -------\n    : :class:`PIL.Image`\n        Image with tracks drawn\n\n    \"\"\"\n    draw = ImageDraw.Draw(image)\n    for track in tracks:\n        bboxes = np.array([np.array(state.state_vector[[0, 2, 4, 5]]).reshape(4)\n                           for state in track.states])\n        x0, y0, w, h = bboxes[-1]\n        x1 = x0 + w\n        y1 = y0 + h\n        draw.rectangle([x0, y0, x1, y1], outline=(255, 0, 0), width=2)\n\n        if show_history:\n            pts = [(box[0] + box[2] / 2, box[1] + box[3] / 2) for box in bboxes]\n            draw.line(pts, fill=(255, 0, 0), width=2)\n\n        class_ = track.metadata['class']['name']\n        score = round(float(track.metadata['score']), 2)\n        if show_class and show_score:\n            draw.text((x0, y1 + 2), '{}:{}'.format(class_, score), fill=(255, 0, 0))\n        elif show_class:\n            draw.text((x0, y1 + 2), '{}'.format(class_), fill=(255, 0, 0))\n        elif show_score:\n            draw.text((x0, y1 + 2), '{}'.format(score), fill=(255, 0, 0))\n    return image\n\n\nfig3, ax3 = plt.subplots(num=\"MultiTargetTracker output\")\nfig3.tight_layout()\nartists3 = []\nprint(\"Running MultiTargetTracker example... Be patient...\")\nfor timestamp, tracks in tracker:\n    if not (len(artists3) + 1) % 10:\n        print(\"Frame: {}/{}\".format(len(artists3) + 1, num_frames))\n\n    # Read the detections\n    detections = detector.detections\n\n    # Read frame\n    frame = frame_reader.frame\n    pixels = copy(frame.pixels)\n\n    # Plot output\n    image = Image.fromarray(pixels)\n    image = draw_detections(image, detections)\n    image = draw_tracks(image, tracks)\n    ax3.axes.xaxis.set_visible(False)\n    ax3.axes.yaxis.set_visible(False)\n    fig3.tight_layout()\n    artist = ax3.imshow(image, animated=True)\n    artists3.append([artist])\nani3 = animation.ArtistAnimation(fig3, artists3, interval=20, blit=True, repeat_delay=200)"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.9.6"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}